import numpy as np
from typing import Dict, List, Optional, Union
import time
from datetime import datetime
import pandas as pd
import json

class ENoseDataGenerator:
    """电子鼻数据生成器"""
    
    def __init__(self, 
                 num_channels: int = 4,
                 sampling_rate: float = 10.0,
                 base_values: Optional[List[float]] = None,
                 noise_level: float = 0.1):
        """
        初始化数据生成器
        
        Args:
            num_channels: 传感器通道数
            sampling_rate: 采样率(Hz)
            base_values: 各通道基准值，如果不指定则随机生成
            noise_level: 噪声水平(0-1之间)
        """
        self.num_channels = num_channels
        self.sampling_rate = sampling_rate
        self.noise_level = noise_level
        
        if base_values is None:
            # 随机生成基准值（1-10之间）
            self.base_values = np.random.uniform(1, 10, num_channels)
        else:
            assert len(base_values) == num_channels, "基准值数量必须等于通道数"
            self.base_values = np.array(base_values)
            
    def generate_sample(self, 
                       num_samples: int,
                       add_drift: bool = True,
                       add_noise: bool = True) -> Dict:
        """
        生成一组传感器数据样本
        
        Args:
            num_samples: 采样点数量
            add_drift: 是否添加漂移
            add_noise: 是否添加噪声
            
        Returns:
            包含传感器数据的字典，符合API要求的格式
        """
        # 生成时间序列
        t = np.linspace(0, num_samples/self.sampling_rate, num_samples)
        
        # 初始化数据矩阵
        data = np.zeros((num_samples, self.num_channels))
        
        for channel in range(self.num_channels):
            # 基准值
            signal = np.ones(num_samples) * self.base_values[channel]
            
            if add_drift:
                # 添加缓慢的漂移（使用正弦函数模拟）
                drift = 0.1 * np.sin(2 * np.pi * 0.01 * t)
                signal += drift
            
            if add_noise:
                # 添加高斯噪声
                noise = np.random.normal(0, self.noise_level, num_samples)
                signal += noise
            
            data[:, channel] = signal
            
        # 生成元数据
        metadata = {
            "temperature": round(np.random.uniform(20, 30), 1),  # 随机温度(20-30℃)
            "humidity": round(np.random.uniform(40, 80), 1),    # 随机湿度(40-80%)
            "pressure": round(np.random.uniform(101.0, 101.5), 3),  # 随机气压
            "timestamp": datetime.now().isoformat(),
            "device_id": "ENose-SIM-001"
        }
        
        return {
            "sensor_readings": data.tolist(),
            "sampling_rate": self.sampling_rate,
            "metadata": metadata
        }
        
    def generate_batch_samples(self, 
                             num_samples: int,
                             batch_size: int = 1,
                             add_drift: bool = True,
                             add_noise: bool = True) -> List[Dict]:
        """
        生成多组传感器数据样本
        
        Args:
            num_samples: 每组的采样点数量
            batch_size: 生成的样本组数
            add_drift: 是否添加漂移
            add_noise: 是否添加噪声
            
        Returns:
            包含多组传感器数据的列表
        """
        return [
            self.generate_sample(num_samples, add_drift, add_noise)
            for _ in range(batch_size)
        ]

    def to_dataframe(self, data: Union[Dict, List[Dict]], include_metadata: bool = True) -> pd.DataFrame:
        """
        将生成的数据转换为DataFrame格式
        
        Args:
            data: 单个或多个数据样本
            include_metadata: 是否包含元数据
            
        Returns:
            pandas DataFrame对象
        """
        if isinstance(data, dict):
            samples = [data]
        else:
            samples = data
            
        # 创建数据列表
        rows = []
        for sample_idx, sample in enumerate(samples):
            readings = sample["sensor_readings"]
            metadata = sample["metadata"]
            
            for time_idx, sensor_values in enumerate(readings):
                row = {
                    "sample_id": sample_idx + 1,
                    "time_point": time_idx + 1,
                    "timestamp": metadata["timestamp"],
                    "time_seconds": time_idx / sample["sampling_rate"]
                }
                
                # 添加传感器数据
                for channel_idx, value in enumerate(sensor_values):
                    row[f"sensor_{channel_idx + 1}"] = value
                    
                # 添加元数据
                if include_metadata:
                    for key, value in metadata.items():
                        if key != "timestamp":  # timestamp已经添加过了
                            row[key] = value
                            
                rows.append(row)
                
        # 创建DataFrame
        df = pd.DataFrame(rows)
        
        # 设置列的顺序
        sensor_cols = [f"sensor_{i+1}" for i in range(self.num_channels)]
        base_cols = ["sample_id", "time_point", "timestamp", "time_seconds"] + sensor_cols
        
        if include_metadata:
            metadata_cols = [col for col in df.columns if col not in base_cols]
            columns = base_cols + metadata_cols
        else:
            columns = base_cols
            
        return df[columns]

    def save_to_file(self, 
                    data: Union[Dict, List[Dict]], 
                    filename: str,
                    include_metadata: bool = True):
        """
        将数据保存到文件
        
        Args:
            data: 要保存的数据
            filename: 输出文件名
            include_metadata: 是否包含元数据
        """
        df = self.to_dataframe(data, include_metadata)
        
        # 根据文件扩展名选择保存格式
        ext = filename.lower().split('.')[-1]
        if ext == 'csv':
            df.to_csv(filename, index=False, encoding='utf-8')
        elif ext in ['xlsx', 'xls']:
            df.to_excel(filename, index=False, engine='openpyxl')
        elif ext == 'json':
            if isinstance(data, dict):
                json_data = data
            else:
                json_data = data
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(json_data, f, indent=2, ensure_ascii=False)
        else:
            raise ValueError(f"不支持的文件格式: {ext}") 